// "Copyright [2021] <Copyright Shengkai Lin>"

#include <memory>
#include "route.h"

using namespace omnetpp;

namespace ai_sim {

bool OxcClusterRouteAlgorithm::isOxcLink(
    TopoManager* topos, const Topology::LinkOut* link, Topology& oxc_topo) {
    topoNode* nodeOxctopo = topos->getNodeinTopo(
        link->getLocalNode(), &oxc_topo);
    topoNode* remoteNodeOxctopo = topos->getNodeinTopo(
        link->getRemoteNode(), &oxc_topo);
    return (nodeOxctopo != nullptr || remoteNodeOxctopo != nullptr);
}

// Flow and gate graph of the optical switch
//
// --------> ------- >---------
//         A |     | B
//         D |     | C
// --------< ------- <----------
//
void OxcClusterRouteAlgorithm::generateRoutes(
    taskboard& tmp_taskboard, TopoManager * topos) {
    // The algorihtm will find the routes of the flows from the oxc
    // and so the time complexity will be reduced
    Topology oxc_topo;
    oxc_topo.extractByModulePath(cStringTokenizer("**.oxc").asVector());
    int stage = 0;
    for (stageboard& routing_stage : tmp_taskboard.stage_info) {
        // Loop in different stages

        // Stores the links that can be reused between different stages
        // include nvlinks and fibers that can be reused between different
        // stages
        vector<topoLink*> stage_tmp_links;
        EV << "***************************stage: " << stage++ << std::endl;
        for (comm_pair& comm_pair : routing_stage.comm_pairs) {
            topos->getDataplaneTopo()->calculateWBSSPathsTo(
                comm_pair.second);
            topoNode* tmp_node = comm_pair.first;

            int first_id = topos->getWorkersTopo()->getNodeId(tmp_node);
            int second_id = topos->getWorkersTopo()->getNodeId(comm_pair.second);
            int gate_id_A = -1, gate_id_B = -1, gate_id_C = -1, gate_id_D = -1;

            EV << "--------------" << "\n";
            EV << "Src is " <<
                tmp_node->getModule()->getFullPath() << "\n";
            EV << "Dest is " << comm_pair.second->getModule()->
                getFullPath() << std::endl;
            EV << "Dist is " << tmp_node->getDistanceToTarget() << std::endl;
            if (tmp_node->getNumPaths() == 0) {
                throw omnetpp::cRuntimeError(
                    "CAN NOT ALLOC ANY ROUTES, CHECK THE TASK");
                return;
            }

            EV << "The src idx" << first_id << endl;
            EV << "The dest idx " << second_id << endl;
            // First, set the output port of the first gpu.
            for (Flow& flow : routing_stage.flows_on_gpu[first_id]) {
                EV << "Flow's src " << flow.srcIdx << endl;
                EV << "Flow's dest " << flow.destIdx << endl;
                if (flow.destIdx == second_id) {
                    flow.outportId = tmp_node->getoutPathtoDest(0)->getLocalGateId();
                    EV << "outport: "<< flow.outportId << endl;
                    EV << "path" << tmp_node->getModule()->gate(flow.outportId)->getFullPath() << endl;
                }
            }
            // Second, Start alloc routes for switches

            // Disable the specified links
            tmp_node->getoutPathtoDest(0)->disable();
            tmp_node->getinPathfromDest(0)->disable();
            // Record the specfied links
            routing_stage.used_links.push_back(tmp_node->getoutPathtoDest(0));
            routing_stage.used_links.push_back(tmp_node->getinPathfromDest(0));

            // If the link is not an OXC link, then we store them in a vector and
            // should release them after the allocation of the stage
            if (!isOxcLink(topos, tmp_node->getoutPathtoDest(0), oxc_topo)) {
                stage_tmp_links.push_back(tmp_node->getoutPathtoDest(0));
                stage_tmp_links.push_back(tmp_node->getinPathfromDest(0));
            }

            gate_id_A = tmp_node->getoutPathtoDest(0)->getRemoteGateId();
            gate_id_D = tmp_node->getinPathfromDest(0)->getRemoteGateId();

            tmp_node = tmp_node->getoutPathtoDest(0)->getRemoteNode();

            while (tmp_node->getDistanceToTarget() > 0) {
                gate_id_B = tmp_node->getoutPathtoDest(0)->getLocalGateId();
                gate_id_C = tmp_node->getinPathfromDest(0)->getLocalGateId();
                RouteInfo tmp_route_f = {
                    .inPort = gate_id_A,
                    .outPort = gate_id_B,
                    .srcIdx = first_id,
                    .destIdx = second_id
                };
                RouteInfo tmp_route_b = {
                    .inPort = gate_id_C,
                    .outPort = gate_id_D,
                    .srcIdx = second_id,
                    .destIdx = first_id
                };
                routing_stage.routes_on_switch[tmp_node->getModuleId()]
                    .push_back(tmp_route_f);
                // Push back depended routes for flows
                for (Flow& flow : routing_stage.flows_on_gpu[first_id]) {
                    if (flow.destIdx == second_id) {
                        flow.dependRoute[tmp_node->getModuleId()].
                            push_back(tmp_route_f);
                    }
                }
                routing_stage.routes_on_switch[tmp_node->getModuleId()]
                    .push_back(tmp_route_b);
                // Push back depended routes for flows
                for (Flow& flow : routing_stage.flows_on_gpu[second_id]) {
                    if (flow.destIdx == first_id) {
                        flow.dependRoute[tmp_node->getModuleId()].
                            push_back(tmp_route_b);
                    }
                }

                EV << "We are in " <<
                    tmp_node->getModule()->getFullPath() << std::endl;
                EV << "Dest is " << comm_pair.second->getModule()->
                    getFullPath() << std::endl;
                EV << tmp_node->getDistanceToTarget() << " hops to go\n";

                tmp_node->getoutPathtoDest(0)->disable();
                tmp_node->getinPathfromDest(0)->disable();

                routing_stage.used_links.push_back(tmp_node->getoutPathtoDest(0));
                routing_stage.used_links.push_back(tmp_node->getinPathfromDest(0));

                // If the link is not an OXC link, then we store them in a vector and
                // should release them after the allocation of the stage
                if (!isOxcLink(topos, tmp_node->getoutPathtoDest(0), oxc_topo)) {
                    stage_tmp_links.push_back(tmp_node->getoutPathtoDest(0));
                    stage_tmp_links.push_back(tmp_node->getinPathfromDest(0));
                }

                gate_id_A = tmp_node->getoutPathtoDest(0)->getRemoteGateId();
                gate_id_D = tmp_node->getinPathfromDest(0)->getRemoteGateId();

                tmp_node = tmp_node->getoutPathtoDest(0)->getRemoteNode();
            }

            // Finally, update the output_port for the backword flows
            for (Flow& flow : routing_stage.flows_on_gpu[second_id]) {
                EV << "Flow's src " << flow.srcIdx << endl;
                EV << "Flow's dest " << flow.destIdx << endl;
                if (flow.destIdx == first_id) {
                    flow.outportId = gate_id_D;
                    EV << "outport: "<< flow.outportId << endl;
                    EV << "path" << tmp_node->getModule()->gate(flow.outportId)->getFullPath() << endl;
                }
            }
        }
        // enable the stage tmp links
        for (topoLink* link : stage_tmp_links) {
            link->enable();
        }
    }
    // Finally disable all the links used by the task
    for (stageboard& stage : tmp_taskboard.stage_info) {
        for (topoLink* link : stage.used_links) {
            link->disable();
        }
    }
}

}  // namespace ai_sim
